import numpy as np
from scipy.interpolate import interp1d
from matplotlib import pyplot as plt
from utils.mcm_algorithm.MCM_common_function import *
import os
import time

# 瑞利测密（适用于DUSM）
def inversion_Rayleigh_density(n_MCM,S_time_sum,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
                               quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,altitudeSystem,
                               altitude_inf,density_inf,uncertain_densityRef):
    # n_MCM：MCM试验次数
    # S_time_sum：探测文件各高度L0级数据RayVHS（求和后），数组[1e8,1e6,……,90]
    # Altitude_final：L2级探测文件数据高度，如[20,21,……,80]
    # dead_time：探测器死区时间
    # uncertain_dead_time：探测器死区时间波动度
    # dark_noise：探测器暗噪声
    # uncertain_dark_noise：探测器暗噪声波动度
    # quantum_efficiency：探测器量子效率
    # uncertain_quantum_efficiency：探测器量子效率波动度
    # pulse_repeat：激光器重频
    # pulse_time：L2级积分时间
    # altitudeSysten_MCM：系统所在高度
    # altitude_inf：参考密度所在高度
    # density_inf：参考密度
    # uncertain_densityRef：参考密度波动度

    # 预处理
    Altitude_final = np.array(Altitude_final).reshape(-1,1).astype(np.float64)
    if np.min(np.abs(Altitude_final-altitude_inf)) > 1:
        return np.full(len(Altitude_final),np.nan)
    height_num = len(S_time_sum)
    S_Noise_inv = preprocessing(n_MCM,S_time_sum,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    # s_std = np.std(S_Noise_inv,axis=1)
    # np.savetxt('s_std.txt', s_std)
    # 密度反演
    ind = np.nanargmin(np.abs(Altitude_final-altitude_inf))
    # density_inf = np.random.normal(density_inf, uncertain_densityRef*0.01*density_inf,[1, n_MCM])
    Density_inv = (Altitude_final-altitudeSystem)**2 * S_Noise_inv / (altitude_inf-altitudeSystem)**2       \
        / S_Noise_inv[ind] * density_inf
    Density_inv = np.abs(Density_inv)
    Density_uncertain = np.std(Density_inv,axis=1)
    
    # Density_mean = Density_inv.mean(axis=1)
    # plt.figure()
    # # plt.semilogx(Density_mean, Altitude_final)
    # plt.plot(Density_mean, Altitude_final)
    # plt.grid(True)
    # np.savetxt('Result/uncer_MCM/RayD_uncer.txt', Density_uncertain)
    # plt.figure()
    # plt.plot(Density_uncertain, Altitude_final)
    # plt.grid(True)

    return Density_uncertain

# 瑞利测温（适用于TUSM）
def inversion_Rayleigh_temperature(n_MCM,S_time_sum,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
                                   quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,altitudeSystem,
                                   altitude_inf,density_inf,uncertain_densityRef,altitude_invStart,Temperature_inf,
                                   uncertain_temperatureRef):
    # n_MCM：MCM试验次数
    # S_time_sum：探测文件各高度L0级数据（求和后）
    # Altitude_final：L2级探测文件数据高度，如[20,21,……,80]
    # dead_time：探测器死区时间
    # uncertain_dead_time：探测器死区时间波动度
    # dark_noise：探测器暗噪声
    # uncertain_dark_noise：探测器暗噪声波动度
    # quantum_efficiency：探测器量子效率
    # uncertain_quantum_efficiency：探测器量子效率波动度
    # pulse_repeat：激光器重频
    # pulse_time：L2级积分时间
    # altitudeSysten_MCM：系统所在高度
    # altitude_inf：参考密度所在高度
    # density_inf：参考密度
    # uncertain_densityRef：参考密度波动度
    # altitude_invStart：密度积分起始高度
    # Temperature_inf：密度积分起始温度
    # uncertain_temperatureRef：参考温度波动度
    
    # 预处理
    Altitude_final = np.array(Altitude_final).reshape(-1,1).astype(np.float64)
    if np.min(np.abs(Altitude_final-altitude_inf)) > 1 or np.min(np.abs(Altitude_final-altitude_invStart)) > 1:
        return np.full(len(Altitude_final),np.nan)
    height_num = len(S_time_sum)
    S_Noise_inv = preprocessing(n_MCM,S_time_sum,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    # 密度反演
    ind = np.nanargmin(np.abs(Altitude_final-altitude_inf))
    # density_inf = np.random.normal(density_inf, uncertain_densityRef*0.01*density_inf,[1, n_MCM])
    Density_inv = (Altitude_final-altitudeSystem)**2 * S_Noise_inv / (altitude_inf-altitudeSystem)**2       \
        / S_Noise_inv[ind] * density_inf
    Density_inv = np.abs(Density_inv)
    # 相关常量
    h = 6.626e-34       # 普朗克常数
    R = 8314.38         # 普适气体常数 J/kmol/K
    NA = 6.022e23       # 阿伏伽德罗常数
    M = 28.9644         # 大气平均分子质量 kg/kmol
    c = 2.998e8         # 光速 m/s
    m = 4.789e-26       # 大气分子平均质量 kg
    G = 6.754e-11       # 引力常数 N*m^2/kg^2
    M_earth = 5.965e24  # 地球质量 kg
    R_earth = 6371.393  # 地球半径 km
    k = 1.38e-23        # 玻尔兹曼常数
    # 温度反演
    Density_inv[Altitude_final.reshape(-1)>altitude_invStart] = np.nan
    g = G * M_earth / (R_earth * 1000 + Altitude_final * 1000) ** 2
    # Temperature_inf = np.random.normal(Temperature_inf,uncertain_temperatureRef*0.01*Temperature_inf,[1, n_MCM])
    Resolution_Ray_final = Altitude_final[1] - Altitude_final[0]                  #单位km
    density_integral = np.nancumsum(Density_inv*g,axis=0) * Resolution_Ray_final     #最低点到z
    ind_temp_ref = np.nanargmin(np.abs(Altitude_final-altitude_invStart))
    density_integral_ref =  density_integral[ind_temp_ref]                        #最低点到z_c
    density_integral_z = density_integral_ref - density_integral                  #z到z_c
    Temperature_inv = Density_inv[ind_temp_ref] / Density_inv * Temperature_inf + M / (R /1000) / Density_inv * density_integral_z
    temp_uncertain = np.std(Temperature_inv,axis=1)

    # Temperature_mean = Temperature_inv.mean(axis=1)
    # plt.figure()
    # plt.plot(Temperature_mean, Altitude_final)
    # plt.grid(True)
    # np.savetxt('Result/uncer_MCM/RayT_uncer.txt', temp_uncertain)
    # plt.figure()
    # plt.plot(temp_uncertain, Altitude_final)
    # plt.grid(True)

    return temp_uncertain

# 拉曼测温（适用于TSVP）
def inversion_Raman(n_MCM,N_sum1,N_sum2,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                    uncertain_quantum_efficiency,pulse_repeat,pulse_time,poly_coef,channel_H=1,channel_L=2):
    # n_MCM：MCM试验次数
    # N_sum1：通道1(RR1)L0级数据（求和后）
    # N_sum2：通道2(RR2)L0级数据（求和后）
    # dead_time：探测器死区时间
    # uncertain_dead_time：探测器死区时间波动度
    # dark_noise：探测器暗噪声
    # uncertain_dark_noise：探测器暗噪声波动度
    # quantum_efficiency：探测器量子效率
    # uncertain_quantum_efficiency：探测器量子效率波动度
    # pulse_repeat：激光器重频
    # pulse_time：L2级积分时间
    # poly_coef：拟合系数，如数组[2,3,1]
    # channel_H：高阶通道编号
    # channel_L：低阶通道编号

    # 预处理
    S_inv_H = preprocessing(n_MCM,N_sum1,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                            uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    S_inv_L = preprocessing(n_MCM,N_sum2,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                            uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    # 温度反演
    Q_Raman = S_inv_H / S_inv_L
    poly_coef = [0.003945,-0.00171,-0.00101]
    a = poly_coef[0]
    b = poly_coef[1]
    c = poly_coef[2]
    Temperature_Raman_1ord_inv2 = 1/(c * np.log(Q_Raman) ** 2 + b * np.log(Q_Raman) + a)
    temp_uncertain = np.std(Temperature_Raman_1ord_inv2,axis=1)
    
    # temp_mean = Temperature_Raman_1ord_inv2.mean(axis=1)
    # plt.figure()
    # plt.plot(temp_mean, 20+np.arange(0,len(N_sum1)))
    # plt.grid(True)
    # plt.figure()
    # plt.plot(Temperature_Raman_1ord_inv2[:,:100], 20+np.arange(0,len(N_sum1)))
    # plt.grid(True)
    # np.savetxt('Result/uncer_MCM/RamT_uncer.txt', temp_uncertain)
    # plt.figure()
    # plt.plot(temp_uncertain, 20+np.arange(0,len(N_sum1)))
    # plt.grid(True)

    return temp_uncertain

# 瑞利测风（适用于HWLS、HWSM，包括经向风速Meridi和纬向风速Zonal）
def inversion_Rayleigh_wind(n_MCM,N_sum1,N_sum2,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
                            quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,wavelength,zenith_angle):
    # n_MCM：MCM试验次数
    # N_sum1：通道1的L0级数据（求和后），对于HWLS经向风速即用RayNLS、纬向风速即用RayWLS ，对于HWSM经向风速即用RayNHS、纬向风速即用RayWHS
    # N_sum2：通道2的L0级数据（求和后），对于HWLS经向风速即用RayNLR、纬向风速即用RayWLR ，对于HWSM经向风速即用RayNHR、纬向风速即用RayWHR
    # Altitude_final：L2级探测文件数据高度，如[20,21,……,80]
    # dead_time：探测器死区时间
    # uncertain_dead_time：探测器死区时间波动度
    # dark_noise：探测器暗噪声
    # uncertain_dark_noise：探测器暗噪声波动度
    # quantum_efficiency：探测器量子效率
    # uncertain_quantum_efficiency：探测器量子效率波动度
    # pulse_repeat：激光器重频
    # pulse_time：L2级积分时间
    # wavelength：激光器中心波长
    # zenith_angle：望远镜天顶角

    # 预处理
    height_num = len(N_sum1)
    Altitude_final = np.array(Altitude_final).reshape(-1,1).astype(np.float64)
    S_Noise_inv_1 = preprocessing(n_MCM,N_sum1,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    S_Noise_inv_2 = preprocessing(n_MCM,N_sum2,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)

    # 响应度计算
    # np.seterr(divide='ignore', invalid='ignore')
    Response = (S_Noise_inv_1 - S_Noise_inv_2) / (S_Noise_inv_1 + S_Noise_inv_2)

    a = function_Response(Response,Altitude_final)

    # 多普勒波长差计算
    wavelength = wavelength*1e-9
    delta_doppler = function_Response(Response,Altitude_final) - wavelength
    # 风速反演
    c = 2.998e8         # 光速 m/s
    Velocity_Los_inv = delta_doppler * c / (2 * wavelength)
    zenith_angle = 90
    Velocity_inv = Velocity_Los_inv/np.sin(zenith_angle * np.pi / 180)
    Velocity_uncertain = np.std(Velocity_inv,axis=1)

    # np.savetxt('Result/uncer_MCM/RayW_uncer.txt', Velocity_uncertain)
    # plt.figure()
    # plt.plot(Velocity_inv .mean(axis=1),Altitude_final)
    # plt.grid(True)
    # plt.show()

    return Velocity_uncertain

# 荧光测风（适用于HWSL）
def inversion_ResW(n_MCM,N_sum1,N_sum2,N_sum3,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
                   quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,zenith_angle):
    # n_MCM：MCM试验次数
    # N_sum1：通道1的L0级数据（求和后），经向风速即用NaNF0、纬向风速即用NaWF0
    # N_sum2：通道2的L0级数据（求和后），经向风速即用NaNF1、纬向风速即用NaWF1
    # N_sum3：通道2的L0级数据（求和后），经向风速即用NaNF2、纬向风速即用NaWF2
    # Altitude_final：L2级探测文件数据高度，如[20,21,……,80]
    # dead_time：探测器死区时间
    # uncertain_dead_time：探测器死区时间波动度
    # dark_noise：探测器暗噪声
    # uncertain_dark_noise：探测器暗噪声波动度
    # quantum_efficiency：探测器量子效率
    # uncertain_quantum_efficiency：探测器量子效率波动度
    # pulse_repeat：激光器重频
    # pulse_time：L2级积分时间
    # zenith_angle：望远镜天顶角
    
    # 预处理
    # T1 = time.time()
    n_MCM = 100   # 加快速度
    height_num = len(N_sum1)
    Altitude_final = np.array(Altitude_final).reshape(-1,1).astype(np.float64)
    S_Noise_inv_1 = preprocessing(n_MCM,N_sum1,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    S_Noise_inv_2 = preprocessing(n_MCM,N_sum2,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    S_Noise_inv_3 = preprocessing(n_MCM,N_sum3,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    # 温风比率计算
    RT = (S_Noise_inv_3 + S_Noise_inv_2) / S_Noise_inv_1
    RW = (S_Noise_inv_3 - S_Noise_inv_2) / S_Noise_inv_1

    # 温风比率拟合函数反演
    [Temperature_inv, Velocity_Los_inv] = function_inv_simplify(RT, RW)
    zenith_angle = 90
    Velocity_inv = Velocity_Los_inv/np.sin(zenith_angle * np.pi / 180)
    Velocity_uncertain = np.std(Velocity_inv,axis=1)

    # T2 = time.time()
    # print('程序运行时间:%s毫秒' % ((T2 - T1)*1000))
    # Velocity_mean = Velocity_inv.mean(axis=1)
    # plt.figure()
    # plt.plot(Velocity_mean, Altitude_final)
    # plt.grid(True)
    # plt.show()

    return Velocity_uncertain

# 荧光测温（适用于TMSL）
def inversion_ResT(n_MCM,N_sum1,N_sum2,N_sum3,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
                   quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,zenith_angle):
    # n_MCM：MCM试验次数
    # N_sum1：通道1的L0级数据（求和后），NaNF0
    # N_sum2：通道2的L0级数据（求和后），NaNF1
    # N_sum3：通道2的L0级数据（求和后），NaNF2
    # Altitude_final：L2级探测文件数据高度，如[20,21,……,80]
    # dead_time：探测器死区时间
    # uncertain_dead_time：探测器死区时间波动度
    # dark_noise：探测器暗噪声
    # uncertain_dark_noise：探测器暗噪声波动度
    # quantum_efficiency：探测器量子效率
    # uncertain_quantum_efficiency：探测器量子效率波动度
    # pulse_repeat：激光器重频
    # pulse_time：L2级积分时间
    # zenith_angle：望远镜天顶角
    
    # 预处理
    # T1 = time.time()
    n_MCM = 100   # 加快速度
    height_num = len(N_sum1)
    Altitude_final = np.array(Altitude_final).reshape(-1,1).astype(np.float64)
    S_Noise_inv_1 = preprocessing(n_MCM,N_sum1,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    S_Noise_inv_2 = preprocessing(n_MCM,N_sum2,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    S_Noise_inv_3 = preprocessing(n_MCM,N_sum3,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
                                uncertain_quantum_efficiency,pulse_repeat,pulse_time)
    # 温风比率计算
    RT = (S_Noise_inv_3 + S_Noise_inv_2) / S_Noise_inv_1
    RW = (S_Noise_inv_3 - S_Noise_inv_2) / S_Noise_inv_1

    # 温风比率拟合函数反演
    [Temperature_inv, Velocity_Los_inv] = function_inv_simplify(RT, RW)
    Temperature_uncertain = np.std(Temperature_inv,axis=1)

    # T2 = time.time()
    # print('程序运行时间:%s毫秒' % ((T2 - T1)*1000))
    # Temperature_mean = Temperature_inv.mean(axis=1)
    # plt.figure()
    # plt.plot(Temperature_mean, Altitude_final)
    # plt.grid(True)
    # plt.show()
    
    return Temperature_uncertain

if __name__ == '__main__':
    n_MCM = 1e6
    n_MCM = int(n_MCM)
    Altitude_final = np.linspace(30,80,61)  # km
    S_time_sum = np.linspace(1e8,1e4,61)
    N_sum1 = np.linspace(2e8,1e4,61)
    N_sum2 = np.linspace(1e8,1e4,61)
    N_sum3 = np.linspace(1e7,1e4,61)
    dead_time = 0.1 # ns
    uncertain_dead_time = 0.1 # %
    dark_noise = 300 # Hz
    uncertain_dark_noise = 0.1 # %
    quantum_efficiency = 10 # %
    uncertain_quantum_efficiency = 0.1 # %
    pulse_repeat = 10 # %
    pulse_time = 1200 # s
    altitude_inf = 40 # km
    density_inf = 0.004052 # kg/m^3
    uncertain_densityRef = 1 # %
    altitudeSystem = 0.5 # km
    altitude_invStart = 80  # km
    Temperature_inf = 186.29  # K
    uncertain_temperatureRef = 1 # %
    poly_coef = [2,3,1] # %
    wavelength = 532  # nm
    zenith_angle = 15   # °

    inversion_Rayleigh_density(n_MCM,S_time_sum,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
                               quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,altitudeSystem,
                               altitude_inf,density_inf,uncertain_densityRef)

    # inversion_Rayleigh_temperature(n_MCM,S_time_sum,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
    #                                quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,altitudeSystem,
    #                                altitude_inf,density_inf,uncertain_densityRef,altitude_invStart,Temperature_inf,
    #                                uncertain_temperatureRef)

    # inversion_Raman(n_MCM,N_sum1,N_sum2,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,quantum_efficiency,
    #                 uncertain_quantum_efficiency,pulse_repeat,pulse_time,poly_coef)

    # inversion_Rayleigh_wind(n_MCM,N_sum1,N_sum2,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
    #                         quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,wavelength,zenith_angle)

    # inversion_ResW(n_MCM,N_sum1,N_sum2,N_sum3,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
    #                quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,zenith_angle)
    
    # inversion_ResT(n_MCM,N_sum1,N_sum2,N_sum3,Altitude_final,dead_time,uncertain_dead_time,dark_noise,uncertain_dark_noise,
    #                quantum_efficiency,uncertain_quantum_efficiency,pulse_repeat,pulse_time,zenith_angle)